package com.mango.data.structure;

import java.util.*;

/**
 * 并查集
 * @Author: mango
 * @Date: 2022/5/22 4:57 下午
 */
public class UnionFindSet<V> {
    // 原始数据对应的节点集合map
    private Map<V,Node<V>> dataMap;
    // 顶层头节点所在的集合大小map
    private Map<Node<V>,Integer> countMap;

    /**
     * 节点
     * @param <V> 原始数据
     */
    private final class Node<V>{
        public V val;
        // 前指针
        public Node<V> pre;

        public Node(V val) {
            this.val = val;
        }

        @Override
        public String toString() {
            return "Node{" +
                    "val=" + val +
                    ", pre.val=" + pre.val +
                    '}';
        }
    }

    public UnionFindSet(List<V> vList){
        dataMap = new HashMap<>();
        countMap = new HashMap<>();
        for(V v : vList){
            Node<V> node = new Node<>(v);
            node.pre = node;
            dataMap.put(v,node);
            countMap.put(node,1);
        }
    }

    /**
     * 判断2个元素是否在同一个集合
     * 找2个头节点，如果头节点相同则说明是同一个集合
     * @param a
     * @param b
     * @return
     */
    public boolean isSameSet(V a, V b){
        return getHead(dataMap.get(a)) == getHead(dataMap.get(b));
    }

    /**
     * 合并2个元素所在的节点到同一个集合
     * 如果是同一个集合，就跳过
     * 找到2个元素的头节点，将小的集合的头节点指向大的集合的头节点
     * 合并后，将原来的节点所在集合数据删除
     * @param a
     * @param b
     */
    public void union(V a, V b){
        if(isSameSet(a, b)){
           return ;
        }
        Node<V> aHead = getHead(dataMap.get(a));
        Node<V> bHead = getHead(dataMap.get(b));
        Integer aCount = countMap.get(aHead);
        Integer bCount = countMap.get(bHead);
        Node<V> bigHead = aCount >= bCount ? aHead : bHead;
        Node<V> smallHead = bigHead == aHead ? bHead : aHead;
        smallHead.pre = bigHead;
        countMap.put(bigHead,aCount + bCount);
        countMap.remove(smallHead);
    }

    /**
     * 通过节点向上找头节点
     * 利用栈来优化节点指向关系为扁平化，加快查询效率为O(1)
     * @param node
     * @return
     */
    private Node<V> getHead(Node<V> node){
        // 使用栈将底层节点都入栈
        Stack<Node<V>> speed = new Stack<>();
        while (!node.val.equals(node.pre.val)){
            node = node.pre;
            speed.push(node);
        }
        // 底层节点出栈，将pre指向head的node节点
        while (!speed.isEmpty()){
            speed.pop().pre = node;
        }
        return node;
    }

    @Override
    public String toString() {
        return "UnionFindSet{" +
                "dataMap=" + dataMap +
                ", countMap=" + countMap +
                '}';
    }

    public static void main(String[] args) {
        List<String> list = new ArrayList<>();
        list.add("A");
        list.add("B");
        list.add("C");
        list.add("D");
        list.add("E");
        UnionFindSet<String> unionFindSet = new UnionFindSet<>(list);
        unionFindSet.union("A","B");
        unionFindSet.union("C","D");
        System.out.println(unionFindSet);
        unionFindSet.union("E","B");
        System.out.println(unionFindSet);
        unionFindSet.union("C","E");
        System.out.println(unionFindSet);
    }
}
